import numpy as np
import gym


def clamp(x, a, b):
    return max(a, min(x, b))


class MyWrapperCartPole(gym.Wrapper):

    def __init__(self, show=False, seed=None):
        render_model = 'human' if show else 'rgb_array'
        env_name = 'CartPole-v0'
        env = gym.make(env_name, render_mode=render_model)
        super().__init__(env)

        if seed:
            env.action_space.seed(seed)
            env.observation_space.seed(seed)
        self.env = env
        self.seed = seed
        self._reset_seed = seed
        self.is_discrete_action = isinstance(env.action_space, gym.spaces.Discrete)

    def reset(self):
        state, _ = self.env.reset(seed=self._reset_seed)
        if self.seed:
            self._reset_seed += 1
        return state

    def step(self, action):
        if isinstance(action, np.ndarray):
            if action.shape[0] != 1:
                action = np.argmax(action)
            else:
                action = action.item()
                action = round(action)
                action = clamp(action, 0, 1)

        res = super().step(action)

        if len(res) == 5:
            state, reward, terminated, truncated, info = res
        else:
            state, reward, terminated, info = res
            truncated = False

        terminated = terminated or truncated

        return state, reward, terminated, info
